#include "Scene.h"

using namespace OBR;

Scene::Scene()
{
}

Scene::~Scene()
{
}

void Scene::setBackground(std::shared_ptr<Background> bg)
{
    this->background = bg;
}

void Scene::setCamera(std::shared_ptr<Camera> cam)
{
    this->camera = cam;
}

void Scene::setIntegrator(std::shared_ptr<Integrator> integrator)
{
    this->integrator = integrator;
}

void Scene::setupLaunchParams()
{
    /* ------------------------ set light desc block ------------------------ */
    int n_lights = this->lights.getInstanceCount();
    LightInfo *light_desc = new LightInfo[n_lights];
    for (int i = 0; i < n_lights; i++)
    {
        std::shared_ptr<Light> l = std::dynamic_pointer_cast<Light>(this->lights[i]);
        light_desc[i] = l->getInfo();
    }
    // upload
    uint memsize = n_lights * sizeof(LightInfo);
    this->d_light_desc.resize(memsize);
    this->d_light_desc.upload(light_desc, memsize);
    // finishing up
    this->params.light_desc = reinterpret_cast<LightInfo *>(this->d_light_desc.d_ptr);
    delete[] light_desc;

    /* --------------------- set callable base offsets ---------------------- */
    this->params.material_base_offset = 1 + n_lights;
    this->params.medium_base_offset =
        this->params.material_base_offset + this->materials.getUniqueCount();
}

void Scene::setupRenderer()
{
    /* -------------------------------- setup AS -------------------------------- */
    this->rootAS->buildAccelStructure(this->renderer.context, this->renderer.stream);

    /* ----------------------------- create pipeline ---------------------------- */
    // setup options
    OptixPipelineCompileOptions compile_options = {};
    compile_options.usesMotionBlur = true;
    compile_options.traversableGraphFlags = OPTIX_TRAVERSABLE_GRAPH_FLAG_ALLOW_ANY;
    compile_options.numPayloadValues = 4;
    compile_options.numAttributeValues = 4;
    compile_options.exceptionFlags = OPTIX_EXCEPTION_FLAG_DEBUG;
    compile_options.pipelineLaunchParamsVariableName = "params";

    // compile and link modules
    this->background->linkProgram(this->renderer.context, compile_options);
    this->camera->linkProgram(this->renderer.context, compile_options);
    this->integrator->linkProgram(this->renderer.context, compile_options);
    for (auto &&c : this->lights.getUnique())
    {
        c->linkProgram(this->renderer.context, compile_options);
    }
    for (auto &&c : this->materials.getUnique())
    {
        c->linkProgram(this->renderer.context, compile_options);
    }
    for (auto &&c : this->mediums.getUnique())
    {
        c->linkProgram(this->renderer.context, compile_options);
    }
    for (auto &&m : this->meshes.getUnique())
    {
        m->linkProgram(this->renderer.context, compile_options);
    }

    // combine
    this->renderer.createPipeline(compile_options);

    /* ------------------------------- create SBT ------------------------------- */
    this->lights.uploadAllPayloads();
    this->materials.uploadAllPayloads();
    this->mediums.uploadAllPayloads();
    this->meshes.uploadAllPayloads();
    this->renderer.createSbt(this->camera, this->background, this->meshes, this->integrator,
                             this->lights, this->materials, this->mediums);
}

void Scene::renderFrame()
{
    this->params.handle = this->rootAS->handle;
    this->renderer.renderFrame(this->params);
    this->postprocess();
    // this->downloadResult();
    /* ------------------------------ update states ----------------------------- */
    this->has_update = false;
    this->frame_idx++;
}

void Scene::postprocess()
{
    if (this->frame_idx == 0 || this->has_update)
    {
        this->d_result.copy(this->renderer.d_fbos);
        this->downloadResult();
        this->acc_sample_cnt = 0;
    }
    else
    {
        /* --------------------------- do gpu post process -------------------------- */
        if (this->enable_accumulate)
        {
            PostProcess::accumulate_radiance_cuda(
                this->params.width, this->params.height, this->renderer.d_fbos.d_pointer(),
                this->d_result.d_pointer(), this->acc_sample_cnt, this->params.samples_per_launch);
        }
        else
        {
            this->acc_sample_cnt = 0;
            this->d_result.copy(this->renderer.d_fbos);
        }
        this->downloadResult();

        /* -------------------------- do cpu post process --------------------------- */
    }
    this->acc_sample_cnt += this->params.samples_per_launch;
}

void Scene::downloadResult()
{
    size_t fbo_size = this->params.width * this->params.height * this->params.num_fbos;
    if (this->result)
    {
        delete[] this->result;
    }
    this->result = new Spectrum[fbo_size];
    this->d_result.download(this->result, fbo_size * sizeof(Spectrum));
}